from flask import Flask, request, jsonify, render_template_string
from langfuse import Langfuse
import requests
import uuid
import json
from datetime import datetime
import logging

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = Flask(__name__)

# 初始化 Langfuse - 根据官方文档配置
langfuse_client = Langfuse(
    public_key="pk-1f-41d47442-ef86e1b-806d-fd8ab170970a",
    secret_key="sk-1f-0a654f31-cd55-4dfe-8896-cd38404ef23e",
    host="https://langfuse-dev.chinabeego.com",
    environment="default",
    debug=True,
    flush_at=1,
    flush_interval=1
)

# 内存存储
conversation_sessions = {}

@app.route('/')
def index():
    """提供前端界面"""
    session_id = str(uuid.uuid4())
    html = """
    <!DOCTYPE html>
    <html lang="zh-CN">
    <head>
        <meta charset="UTF-8">
        <meta name="viewport" content="width=device-width, initial-scale=1.0">
        <title>Langfuse 分数测试应用</title>
        <style>
            * {
                margin: 0;
                padding: 0;
                box-sizing: border-box;
            }
            
            body {
                font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif;
                background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
                min-height: 100vh;
                color: #333;
                padding: 20px;
            }
            
            .container {
                max-width: 800px;
                margin: 0 auto;
                background: white;
                border-radius: 20px;
                box-shadow: 0 20px 40px rgba(0,0,0,0.1);
                overflow: hidden;
            }
            
            .header {
                background: linear-gradient(135deg, #4f46e5 0%, #7c3aed 100%);
                color: white;
                padding: 20px;
                text-align: center;
            }
            
            .header h1 {
                font-size: 2rem;
                margin-bottom: 10px;
            }
            
            .content {
                padding: 20px;
            }
            
            .test-section {
                margin-bottom: 30px;
                padding: 20px;
                border: 1px solid #e2e8f0;
                border-radius: 12px;
                background: #f8fafc;
            }
            
            .test-section h2 {
                margin-bottom: 15px;
                color: #4f46e5;
            }
            
            .test-btn {
                background: linear-gradient(135deg, #3b82f6, #1d4ed8);
                color: white;
                border: none;
                padding: 12px 24px;
                border-radius: 12px;
                font-size: 1rem;
                font-weight: 600;
                cursor: pointer;
                transition: all 0.3s ease;
                margin-right: 10px;
            }
            
            .test-btn:hover {
                transform: translateY(-2px);
                box-shadow: 0 6px 20px rgba(59, 130, 246, 0.4);
            }
            
            .result {
                margin-top: 15px;
                padding: 15px;
                border-radius: 8px;
                background: white;
                border: 1px solid #e2e8f0;
            }
            
            .success {
                color: #059669;
                border-color: #059669;
            }
            
            .error {
                color: #dc2626;
                border-color: #dc2626;
            }
            
            .loading {
                display: none;
                text-align: center;
                padding: 20px;
                color: #64748b;
            }
            
            .loading.show {
                display: block;
            }
            
            .trace-info {
                margin-top: 20px;
                padding: 15px;
                background: #f1f5f9;
                border-radius: 8px;
            }
        </style>
    </head>
    <body>
        <div class="container">
            <div class="header">
                <h1>Langfuse 分数上传测试</h1>
                <p>测试 Langfuse 分数上传功能</p>
            </div>
            
            <div class="content">
                <div class="test-section">
                    <h2>测试分数上传</h2>
                    <p>使用写死的分数值测试 Langfuse 分数上传功能</p>
                    <button class="test-btn" onclick="testScores()">测试分数上传</button>
                    <div class="result" id="scoreResult"></div>
                </div>
                
                <div class="test-section">
                    <h2>检查分数状态</h2>
                    <p>输入 Trace ID 检查分数状态</p>
                    <input type="text" id="traceIdInput" placeholder="输入 Trace ID" style="padding: 10px; width: 300px; margin-right: 10px;">
                    <button class="test-btn" onclick="checkScores()">检查分数</button>
                    <div class="result" id="checkResult"></div>
                </div>
                
                <div class="test-section">
                    <h2>创建完整对话追踪</h2>
                    <p>创建包含对话和分数的完整追踪记录</p>
                    <input type="text" id="messageInput" placeholder="输入测试消息" style="padding: 10px; width: 300px; margin-right: 10px;">
                    <button class="test-btn" onclick="createFullTrace()">创建追踪</button>
                    <div class="result" id="traceResult"></div>
                </div>
                
                <div class="loading" id="loading">
                    处理中，请稍候...
                </div>
            </div>
        </div>
        
        <script>
            function testScores() {
                document.getElementById('loading').classList.add('show');
                document.getElementById('scoreResult').innerHTML = '';
                
                fetch('/test-scores', {
                    method: 'GET'
                })
                .then(response => response.json())
                .then(data => {
                    document.getElementById('loading').classList.remove('show');
                    if (data.status === 'success') {
                        document.getElementById('scoreResult').className = 'result success';
                        document.getElementById('scoreResult').innerHTML = `
                            <strong>成功:</strong> ${data.message}
                            <div class="trace-info">
                                <p>Trace ID: ${data.trace_id}</p>
                                <p>生成的分数: ${data.scores_count}</p>
                                <p>分数 IDs: ${data.score_ids.join(', ')}</p>
                            </div>
                        `;
                    } else {
                        document.getElementById('scoreResult').className = 'result error';
                        document.getElementById('scoreResult').innerHTML = `
                            <strong>错误:</strong> ${data.message}
                            <p>${data.error || ''}</p>
                        `;
                    }
                })
                .catch(error => {
                    document.getElementById('loading').classList.remove('show');
                    document.getElementById('scoreResult').className = 'result error';
                    document.getElementById('scoreResult').innerHTML = `
                        <strong>错误:</strong> 请求失败
                        <p>${error}</p>
                    `;
                });
            }
            
            function checkScores() {
                const traceId = document.getElementById('traceIdInput').value;
                if (!traceId) {
                    alert('请输入 Trace ID');
                    return;
                }
                
                document.getElementById('loading').classList.add('show');
                document.getElementById('checkResult').innerHTML = '';
                
                fetch('/check-scores/' + traceId, {
                    method: 'GET'
                })
                .then(response => response.json())
                .then(data => {
                    document.getElementById('loading').classList.remove('show');
                    if (data.scores_count !== undefined) {
                        document.getElementById('checkResult').className = 'result success';
                        let scoresHtml = '<ul>';
                        if (data.scores && data.scores.length > 0) {
                            data.scores.forEach(score => {
                                scoresHtml += `<li>${score.name}: ${score.value} (ID: ${score.id})</li>`;
                            });
                        } else {
                            scoresHtml += '<li>未找到分数</li>';
                        }
                        scoresHtml += '</ul>';
                        
                        document.getElementById('checkResult').innerHTML = `
                            <strong>检查结果:</strong>
                            <p>Trace ID: ${data.trace_id}</p>
                            <p>分数数量: ${data.scores_count}</p>
                            ${scoresHtml}
                        `;
                    } else {
                        document.getElementById('checkResult').className = 'result error';
                        document.getElementById('checkResult').innerHTML = `
                            <strong>错误:</strong> ${data.error || '未知错误'}
                        `;
                    }
                })
                .catch(error => {
                    document.getElementById('loading').classList.remove('show');
                    document.getElementById('checkResult').className = 'result error';
                    document.getElementById('checkResult').innerHTML = `
                        <strong>错误:</strong> 请求失败
                        <p>${error}</p>
                    `;
                });
            }
            
            function createFullTrace() {
                const message = document.getElementById('messageInput').value;
                if (!message) {
                    alert('请输入测试消息');
                    return;
                }
                
                document.getElementById('loading').classList.add('show');
                document.getElementById('traceResult').innerHTML = '';
                
                fetch('/create-trace', {
                    method: 'POST',
                    headers: {
                        'Content-Type': 'application/json'
                    },
                    body: JSON.stringify({
                        message: message,
                        session_id: 'test-session-' + Date.now()
                    })
                })
                .then(response => response.json())
                .then(data => {
                    document.getElementById('loading').classList.remove('show');
                    if (data.status === 'success') {
                        document.getElementById('traceResult').className = 'result success';
                        document.getElementById('traceResult').innerHTML = `
                            <strong>成功:</strong> ${data.message}
                            <div class="trace-info">
                                <p>Trace ID: ${data.trace_id}</p>
                                <p>生成的分数: ${data.scores_count}</p>
                                <p>AI 回复: ${data.response}</p>
                            </div>
                        `;
                    } else {
                        document.getElementById('traceResult').className = 'result error';
                        document.getElementById('traceResult').innerHTML = `
                            <strong>错误:</strong> ${data.message}
                            <p>${data.error || ''}</p>
                        `;
                    }
                })
                .catch(error => {
                    document.getElementById('loading').classList.remove('show');
                    document.getElementById('traceResult').className = 'result error';
                    document.getElementById('traceResult').innerHTML = `
                        <strong>错误:</strong> 请求失败
                        <p>${error}</p>
                    `;
                });
            }
        </script>
    </body>
    </html>
    """
    return render_template_string(html, session_id=session_id)

@app.route('/test-scores')
def test_scores():
    """测试分数上传功能的专用端点"""
    try:
        # 创建一个测试 trace
        trace = langfuse_client.trace(
            name="test_score_trace",
            input={"test": "score_test"},
            metadata={"test": True}
        )
        
        # 创建一个测试 generation
        generation = trace.generation(
            name="test_generation",
            input="Test input for scoring",
            output="Test output for scoring"
        )
        
        logger.info(f"创建测试 trace: {trace.id}")
        logger.info(f"创建测试 generation: {generation.id}")
        
        # 使用 create_score 方法创建分数
        score_ids = []
        
        # 方法1: 使用 create_score 方法
        score1 = langfuse_client.create_score(
            name="test_score_1",
            value=0.95,  # 写死的分数值
            trace_id=trace.id,
            observation_id=generation.id,
            comment="测试分数 1"
        )
        score_ids.append(score1.id)
        
        # 方法2: 使用 create_score 方法
        score2 = langfuse_client.create_score(
            name="test_score_2",
            value=0.85,  # 写死的分数值
            trace_id=trace.id,
            observation_id=generation.id,
            comment="测试分数 2"
        )
        score_ids.append(score2.id)
        
        # 方法3: 使用 create_score 方法
        score3 = langfuse_client.create_score(
            name="test_score_3",
            value=0.75,  # 写死的分数值
            trace_id=trace.id,
            observation_id=generation.id,
            comment="测试分数 3"
        )
        score_ids.append(score3.id)
        
        # 强制刷新
        langfuse_client.flush()
        
        return jsonify({
            "status": "success",
            "message": "测试分数创建成功",
            "trace_id": trace.id,
            "generation_id": generation.id,
            "scores_count": 3,
            "score_ids": score_ids
        })
        
    except Exception as e:
        logger.error(f"测试分数创建失败: {e}")
        import traceback
        logger.error(traceback.format_exc())
        return jsonify({
            "status": "error",
            "message": f"测试分数创建失败: {str(e)}",
            "error": str(e)
        }), 500

@app.route('/check-scores/<trace_id>')
def check_scores(trace_id):
    """检查特定 trace 的分数状态"""
    try:
        # 获取该 trace 的所有分数
        scores = langfuse_client.get_scores(trace_id=trace_id)
        
        return jsonify({
            "trace_id": trace_id,
            "scores_count": len(scores) if scores else 0,
            "scores": [{
                "id": score.id,
                "name": score.name,
                "value": score.value,
                "observation_id": score.observation_id,
                "comment": score.comment,
                "timestamp": score.timestamp.isoformat() if hasattr(score, 'timestamp') else None
            } for score in scores] if scores else []
        })
    except Exception as e:
        return jsonify({
            "error": str(e),
            "trace_id": trace_id
        }), 500

@app.route('/create-trace', methods=['POST'])
def create_trace():
    """创建包含对话和分数的完整追踪"""
    data = request.get_json()
    user_message = data['message']
    session_id = data['session_id']
    
    try:
        # 创建 trace
        trace = langfuse_client.trace(
            name="dify_chat",
            input={"message": user_message, "session_id": session_id},
            metadata={"provider": "dify", "test": True}
        )
        
        # 创建 generation
        generation = trace.generation(
            name="dify_api_call",
            model="dify-ai-agent",
            input=user_message,
            metadata={"session_id": session_id}
        )
        
        # 模拟 AI 回复
        assistant_reply = f"这是对 '{user_message}' 的模拟回复"
        
        # 更新 generation
        generation.end(
            output=assistant_reply,
            usage={
                "prompt_tokens": len(user_message) // 4,
                "completion_tokens": len(assistant_reply) // 4,
                "total_tokens": (len(user_message) + len(assistant_reply)) // 4
            },
            metadata={"simulated": True}
        )
        
        # 创建多个测试分数
        score_ids = []
        
        # 创建质量分数
        score1 = langfuse_client.create_score(
            name="quality_score",
            value=0.8,  # 写死的分数值
            trace_id=trace.id,
            observation_id=generation.id,
            comment="质量评分"
        )
        score_ids.append(score1.id)
        
        # 创建相关性分数
        score2 = langfuse_client.create_score(
            name="relevance_score",
            value=0.9,  # 写死的分数值
            trace_id=trace.id,
            observation_id=generation.id,
            comment="相关性评分"
        )
        score_ids.append(score2.id)
        
        # 创建有用性分数
        score3 = langfuse_client.create_score(
            name="helpfulness_score",
            value=0.7,  # 写死的分数值
            trace_id=trace.id,
            observation_id=generation.id,
            comment="有用性评分"
        )
        score_ids.append(score3.id)
        
        # 强制刷新
        langfuse_client.flush()
        
        return jsonify({
            "status": "success",
            "message": "追踪和分数创建成功",
            "trace_id": trace.id,
            "response": assistant_reply,
            "scores_count": 3,
            "score_ids": score_ids
        })
        
    except Exception as e:
        logger.error(f"创建追踪失败: {e}")
        import traceback
        logger.error(traceback.format_exc())
        return jsonify({
            "status": "error",
            "message": f"创建追踪失败: {str(e)}",
            "error": str(e)
        }), 500

if __name__ == '__main__':
    app.run(debug=True, host='0.0.0.0', port=5001)